#include <zenoui/ImageIO.hpp>
#include <cstring>
#include <string>
#include <vector>
#include <stdexcept>
#include <zenox/minilog.hpp>
#include <zenox/unicode.hpp>
#include <zenox/glm.hpp>
#define STB_IMAGE_STATIC
#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_STATIC
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <zenox/push_warn.hpp>
#include <zeno3rd/stb_image.h>
#include <zeno3rd/stb_image_write.h>
#include <zenox/pop_warn.hpp>

namespace zenoui {

template <>
void Image<glm::u8vec4>::load(std::string const &path) {
    int nx = 0, ny = 0, comp = 4;
    unsigned char *p = stbi_load(zenox::u8path(path).string().c_str(), &nx, &ny, &comp, comp);
    if (comp != 4) [[unlikely]] throw std::invalid_argument("invalid component return from stbi_load");
    m_flatimg.resize(nx * ny);
    std::memcpy(m_flatimg.data(), p, m_flatimg.size() * sizeof(m_flatimg[0]));
    stbi_image_free(p);
    m_nx = nx;
    m_ny = ny;
}

template <>
void Image<glm::u8vec3>::load(std::string const &path) {
    int nx = 0, ny = 0, comp = 3;
    unsigned char *p = stbi_load(zenox::u8path(path).string().c_str(), &nx, &ny, &comp, comp);
    if (comp != 3) [[unlikely]] throw std::invalid_argument("invalid component return from stbi_load");
    m_flatimg.resize(nx * ny);
    std::memcpy(m_flatimg.data(), p, m_flatimg.size() * sizeof(m_flatimg[0]));
    stbi_image_free(p);
    m_nx = nx;
    m_ny = ny;
}

template <>
void Image<std::uint8_t>::load(std::string const &path) {
    int nx = 0, ny = 0, comp = 1;
    unsigned char *p = stbi_load(zenox::u8path(path).string().c_str(), &nx, &ny, &comp, comp);
    if (comp != 1) [[unlikely]] throw std::invalid_argument("invalid component return from stbi_load");
    m_flatimg.resize(nx * ny);
    std::memcpy(m_flatimg.data(), p, m_flatimg.size() * sizeof(m_flatimg[0]));
    stbi_image_free(p);
    m_nx = nx;
    m_ny = ny;
}

template <>
void Image<glm::vec4>::load(std::string const &path) {
    int nx = 0, ny = 0, comp = 4;
    float *p = stbi_loadf(zenox::u8path(path).string().c_str(), &nx, &ny, &comp, comp);
    if (comp != 4) [[unlikely]] throw std::invalid_argument("invalid component return from stbi_loadf");
    m_flatimg.resize(nx * ny);
    std::memcpy(m_flatimg.data(), p, m_flatimg.size() * sizeof(m_flatimg[0]));
    stbi_image_free(p);
    m_nx = nx;
    m_ny = ny;
}

static std::vector<float> u8tof(const void *data, size_t n) { // 0~255 to 0.0~1.0
    auto datau8 = static_cast<const std::uint8_t *>(data);
    std::vector<float> result(n);
    for (size_t i = 0; i < n; ++i) {
        result[i] = std::pow(static_cast<float>(datau8[i]) / 255.0f, stbi__l2h_gamma) * stbi__l2h_scale;
    }
    return result;
}

static std::vector<std::uint8_t> ftou8(const float *data, size_t n) { // 0.0~1.0 to 0~255
    std::vector<std::uint8_t> result(n);
    for (size_t i = 0; i < n; ++i) {
        float z = std::pow(data[i] * stbi__h2l_scale_i, stbi__h2l_gamma_i) * 255 + 0.5f;
        if (z < 0) z = 0;
        if (z > 255) z = 255;
        result[i] = (stbi_uc) stbi__float2int(z);
    }
    return result;
}

static void stbi_write_generic(char const *filename, int x, int y, int comp, const void *data, std::string const &type, int quality) {
    auto nativefname = zenox::u8path(filename).string();
    auto filename1 = nativefname.c_str();
    int ret;
    if (type == "png") {
        ret = stbi_write_png(filename1, x, y, comp, data, 0);
    } else if (type == "bmp") {
        ret = stbi_write_bmp(filename1, x, y, comp, data);
    } else if (type == "jpg") {
        ret = stbi_write_jpg(filename1, x, y, comp, data, quality);
    } else if (type == "tga") {
        ret = stbi_write_tga(filename1, x, y, comp, data);
    } else if (type == "hdr") {
        auto fdata = u8tof(data, x * y * comp);
        ret = stbi_write_hdr(filename1, x, y, comp, fdata.data());
    } else [[unlikely]] {
        throw std::invalid_argument("invalid extension name for image: [" + type + ']');
    }
    if (!ret) {
        zenox::fatal_error("failed to open file '{}': {}", filename, std::strerror(errno));
    }
}

static void stbi_write_genericf(char const *filename, int x, int y, int comp, const float *data, std::string const &type, int quality) {
    auto nativefname = zenox::u8path(filename).string();
    auto filename1 = nativefname.c_str();
    int ret;
    if (type == "hdr") {
        ret = stbi_write_hdr(filename1, x, y, comp, data);
    } else if (type == "png") {
        auto udata = ftou8(data, x * y * comp);
        ret = stbi_write_png(filename1, x, y, comp, udata.data(), 0);
    } else if (type == "bmp") {
        auto udata = ftou8(data, x * y * comp);
        ret = stbi_write_bmp(filename1, x, y, comp, udata.data());
    } else if (type == "jpg") {
        auto udata = ftou8(data, x * y * comp);
        ret = stbi_write_jpg(filename1, x, y, comp, udata.data(), quality);
    } else if (type == "tga") {
        auto udata = ftou8(data, x * y * comp);
        ret = stbi_write_tga(filename1, x, y, comp, udata.data());
    } else [[unlikely]] {
        throw std::invalid_argument("invalid extension name for image: [" + type + ']');
    }
    if (!ret) {
        zenox::fatal_error("failed to open file '{}': {}", filename, std::strerror(errno));
    }
}

template <>
void Image<glm::vec3>::load(std::string const &path) {
    int nx = 0, ny = 0, comp = 3;
    float *p = stbi_loadf(path.c_str(), &nx, &ny, &comp, comp);
    if (comp != 3) [[unlikely]] throw std::invalid_argument("invalid component return from stbi_loadf");
    m_flatimg.resize(nx * ny);
    std::memcpy(m_flatimg.data(), p, m_flatimg.size() * sizeof(m_flatimg[0]));
    stbi_image_free(p);
    m_nx = nx;
    m_ny = ny;
}

template <>
void Image<float>::load(std::string const &path) {
    int nx = 0, ny = 0, comp = 1;
    float *p = stbi_loadf(path.c_str(), &nx, &ny, &comp, comp);
    if (comp != 1) [[unlikely]] throw std::invalid_argument("invalid component return from stbi_loadf");
    m_flatimg.resize(nx * ny);
    std::memcpy(m_flatimg.data(), p, m_flatimg.size() * sizeof(m_flatimg[0]));
    stbi_image_free(p);
    m_nx = nx;
    m_ny = ny;
}

template <>
void Image<glm::u8vec4>::save(std::string const &path, std::string const &type, int quality) {
    stbi_write_generic(path.c_str(), m_nx, m_ny, 4, m_flatimg.data(), type, quality);
}

template <>
void Image<glm::u8vec3>::save(std::string const &path, std::string const &type, int quality) {
    stbi_write_generic(path.c_str(), m_nx, m_ny, 3, m_flatimg.data(), type, quality);
}

template <>
void Image<std::uint8_t>::save(std::string const &path, std::string const &type, int quality) {
    stbi_write_generic(path.c_str(), m_nx, m_ny, 1, m_flatimg.data(), type, quality);
}

template <>
void Image<glm::vec4>::save(std::string const &path, std::string const &type, int quality) {
    stbi_write_genericf(path.c_str(), m_nx, m_ny, 4, (const float *)m_flatimg.data(), type, quality);
}

template <>
void Image<glm::vec3>::save(std::string const &path, std::string const &type, int quality) {
    stbi_write_genericf(path.c_str(), m_nx, m_ny, 3, (const float *)m_flatimg.data(), type, quality);
}

template <>
void Image<float>::save(std::string const &path, std::string const &type, int quality) {
    stbi_write_genericf(path.c_str(), m_nx, m_ny, 1, m_flatimg.data(), type, quality);
}

}
